/***************************************************************************\
 * The Mattise Kernel														*
 * Copyright 2007 Matthew Iselin											*
 * Licensed under the GPL													*
 *																			*
 * reloc_elf.cc																*
 *																			*
 * Performs relocation on an ELF executable									*
 *																			*
\***************************************************************************/

// Includes
#include <system.h>
#include <stdio.h>
#include <reloc.h>
#include <reloc_elf.h>
#include <memory.h>
#include <keyboard.h>
#include <string.h>
#include <task.h>

// Class Definition

// finds the offset of a table of some description
int ELFRelocation::FindOffset( Elf32_Shdr* __sects, int cnt, int type, int notval )
{
	// find it
	for( int i = 0; i < cnt; i++ )
		if( __sects[i].sh_type == type && i != notval )
			return i;
	return 0;
}
int ELFRelocation::FindOffset( Elf32_Shdr* __sects, int cnt, const char* secname, Elf32_Ehdr* myHead )
{
	// find the section name string table
	char* SHStringTable = (char*) (__sects[myHead->e_shstrndx].sh_offset + GetBase());
	
	// loop through each section
	for( int i = 0; i < cnt; i++ )
	{
		if( strcmp( secname, (char*) &SHStringTable[__sects[i].sh_name+1], strlen( secname ) - 1 ) )
		{
			return i;
		}
	}
	return -1;
}

bool ELFRelocation::RelocateExecutable()
{	
	// elf file header
	Elf32_Ehdr* myHead = (Elf32_Ehdr*) GetBase();
	
	// get the base of the file
	uint_t myFileBase = GetBase();
	
	// object file, or executable?
	if( myHead->e_type == ET_EXEC )
	{
		// executable file (ie. multiple object files linked)
		
		/** get the program header table **/
		
		// the program header table
		Elf32_Phdr* myProgramHeaders = (Elf32_Phdr*) (myHead->e_phoff + myFileBase);
		
		// print how many there are
/*		kprintf( "%d Program Headers\n%d Section Headers\n", myHead->e_phnum, myHead->e_shnum );
		
		// iterate through each one
		int i,j;
		for( i = 0; i < myHead->e_phnum; i++ )
		{
			// tell me info!
			kprintf( "PHDR[%d].p_type = 0x%x\n", i, myProgramHeaders[i].p_type );
			kprintf( "PHDR[%d].p_offset = 0x%x\n", i, myProgramHeaders[i].p_offset );
			kprintf( "PHDR[%d].p_vaddr = 0x%x\n", i, myProgramHeaders[i].p_vaddr );
			kprintf( "PHDR[%d].p_paddr = 0x%x\n", i, myProgramHeaders[i].p_paddr );
			kprintf( "PHDR[%d].p_filesz = 0x%x\n", i, myProgramHeaders[i].p_filesz );
			kprintf( "PHDR[%d].p_memsz = 0x%x\n", i, myProgramHeaders[i].p_memsz );
			kprintf( "PHDR[%d].p_flags = 0x%x\n", i, myProgramHeaders[i].p_flags );
			kprintf( "PHDR[%d].p_align = 0x%x\n", i, myProgramHeaders[i].p_align );
		}*/
	}
	//else if( myHead->e_type == ET_REL )
	{
		// relocatable file (ie. object file)
	
		/** get ready for relocation **/

		// all the section headers
		Elf32_Shdr* mySectionHeaders = (Elf32_Shdr*) (myHead->e_shoff + myFileBase);
		
		// setup the section header string table
		char* SHStringTable = (char*) (mySectionHeaders[myHead->e_shstrndx].sh_offset + GetBase());
		
		// symbol tables
		Elf32_Sym** mySymbolTables;
		int mySymTabCnt = 0;
		
		// string tables (not the section header one, which always exists)
		char** myStrTabs;
		int myStrTabCnt = 0;
		
		// relocation sections
		Elf32_Rel** myRelocation;
		int myRelCnt = 0;
		int myRelSH[32];
		Elf32_Rela** myRelocationA;
		int myRelCntA = 0;
		int myRelASH[32];
		
		// loop indexes
		int i, j;
		
		/** find string tables, symbol tables, and relocation entries **/
		
		// read each section header
		for( i = 0; i < myHead->e_shnum; i++ )
		{
			// what type is this section header?
			switch( mySectionHeaders[i].sh_type )
			{
				// symbol table
				case SHT_SYMTAB:
				
					// read in the table
					mySymbolTables[mySymTabCnt++] = (Elf32_Sym*) (mySectionHeaders[i].sh_offset + myFileBase);
					
					break;
				
				// string table
				case SHT_STRTAB:
				
					// read it in, if it's not the section header string table
					if( i != myHead->e_shstrndx )
						myStrTabs[myStrTabCnt++] = (char*) (mySectionHeaders[i].sh_offset + myFileBase);
					
					break;
				
				// relocation
				case SHT_REL:
				
					// read it in
					myRelSH[myRelCnt] = i;
					myRelocation[myRelCnt++] = (Elf32_Rel*) (mySectionHeaders[i].sh_offset + myFileBase);
					
					break;
					
				case SHT_RELA:
				
					// read it in
					myRelASH[myRelCnt] = i;
					myRelocationA[myRelCntA++] = (Elf32_Rela*) (mySectionHeaders[i].sh_offset + myFileBase);
				
					break;
			}
		}
		
		/** relocate the program **/
		
		// work with each of the relocation sections
		for( i = 0; i < myRelCnt; i++ )
		{
			// the section header referred to here
			Elf32_Shdr thisRelSec = mySectionHeaders[ myRelSH[i] ];
			
			// our runing relocation information array
			Elf32_Rel* reloc = (Elf32_Rel*) ( thisRelSec.sh_offset + myFileBase );
			
			// loop through each entry
			for( j = 0; j < ( thisRelSec.sh_size / thisRelSec.sh_entsize ); j++ )
			{
				// get the offset of the symbol
				int symOff = ELF32_R_SYM( reloc[j].r_info );
				
				// section in which data needs to be changed
				Elf32_Shdr thisChgSec = mySectionHeaders[thisRelSec.sh_info];
				
				// calculate the pointer to the data to change
				uint_t* datOffset = (uint_t*) ( thisChgSec.sh_offset + reloc[j].r_offset + myFileBase );
				
				// get the symbol table this one is referring to
				Elf32_Sym* thisSymTab = (Elf32_Sym*) ( mySectionHeaders[thisRelSec.sh_link].sh_offset + myFileBase );
				
				// read in any addend (this is not a SHT_RELA section)
				uint_t addend = *datOffset;
				
				// set the new data
				*datOffset = mySectionHeaders[thisSymTab[symOff].st_shndx].sh_offset + thisSymTab[symOff].st_value + myFileBase + addend;
			}
		}
		for( i = 0; i < myRelCntA; i++ )
		{
			// the section header referred to here
			Elf32_Shdr thisRelSec = mySectionHeaders[ myRelASH[i] ];
			
			// our runing relocation information array
			Elf32_Rela* reloc = (Elf32_Rela*) ( thisRelSec.sh_offset + myFileBase );
			
			// loop through each entry
			for( j = 0; j < ( thisRelSec.sh_size / thisRelSec.sh_entsize ); j++ )
			{
				// get the offset of the symbol
				int symOff = ELF32_R_SYM( reloc[j].r_info );
				
				// section in which data needs to be changed
				Elf32_Shdr thisChgSec = mySectionHeaders[thisRelSec.sh_info];
				
				// calculate the pointer to the data to change
				uint_t* datOffset = (uint_t*) ( thisChgSec.sh_offset + reloc[j].r_offset + myFileBase );
				
				// get the symbol table this one is referring to
				Elf32_Sym* thisSymTab = (Elf32_Sym*) ( mySectionHeaders[thisRelSec.sh_link].sh_offset + myFileBase );
				
				// read in any addend (this is not a SHT_RELA section)
				uint_t addend = reloc[j].r_addend;
				
				// set the new data
				*datOffset = mySectionHeaders[thisSymTab[symOff].st_shndx].sh_offset + thisSymTab[symOff].st_value + myFileBase + addend;
			}
		}
		
		/** find the first executable section and execute it **/
		
		// executed section index
		int execIndex = -1;

		// search the sections for executable instructions
		for( i = 0; i < myHead->e_shnum; i++ )
		{
			// what are the attributes?
			if( mySectionHeaders[i].sh_flags == ( SHF_ALLOC + SHF_EXECINSTR ) )
			{
				// make sure it's a valid type
				if( mySectionHeaders[i].sh_type == SHT_PROGBITS )
				{
					// this is the one
					execIndex = i;
					
					// don't search anymore
					break;
				}
			}
		}
		
		// check for a valid offset
		if( execIndex == -1 )
			return false; // no executable sections!
		
		/** execution! **/
		
		// execute the program
		// TODO: use a different method
		//CreateThread( mySectionHeaders[execIndex].sh_offset + myFileBase, 1, fileName, 10 );
		
		// return success
		return true;
	}
	
	// no success...
	return false;
}
